#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "rsa.h"
#include "auxl.h"
#include "calc.h"
#include "prim.h"

void generate_keys(big_integer *n, big_integer *e, big_integer *d) {
	big_integer p, q, t;
	while (1) {
		get_random_prime(&p, K_BITS / 2);
		printf("found p: ");
		print_big_integer(&p);
		get_random_prime(&q, K_BITS / 2);
		printf("found q: ");
		print_big_integer(&q);

		if (bi_cmp(&p, &q) > 0) {
			bi_minus(&t, &p, &q);
		}
		else {
			bi_minus(&t, &q, &p);
		}
		int bits = get_bits(&t);
		if (bits <= K_BYTES / 2 - 100) {
			printf("log|p - q| = %d, too small!\n", bits);
			continue;
		}

		bi_mul(n, &p, &q);
		bits = get_bits(n);
		if (bits != K_BITS) {
			printf("n has %d bits, not %d\n", bits, K_BITS);
			continue;
		}

		bits = get_naf_weight(n);
		if (bits <= K_BITS / 4) {
			printf("the naf weight of n is %d, too small!\n", bits);
			continue;
		}

		clear_big_integer(e);
		e->digits[0] = 0x10001;
		big_integer s;
		clear_big_integer(&s);
		s.digits[0] = 1;
		bi_minus(&p, &p, &s);
		bi_minus(&q, &q, &s);
		bi_mul(&t, &p, &q);
		get_inverse(d, e, &t);
		bits = get_bits(d);
		if (bits <= K_BITS / 2) {
			printf("d has %d bits, too small!\n", bits);
			continue;
		}

		break;
	}
}

int rsaes_pkcs1_v15_encrypt(octet_string *C, const big_integer *n, const big_integer *e, const uint8_t *M, size_t mlen) {
	static octet_string em;
	static big_integer m, c;
	if (mlen > K_BYTES - 11) {
		return 1;
	}
	em.octets[0] = 0x00;
	em.octets[1] = 0x02;
	for (int i = 0; i < K_BYTES - mlen - 3; ++i) {
		do {
			em.octets[2 + i] = (rand() & 0xFF);
		} while (em.octets[2 + i] == 0);
	}
	em.octets[K_BYTES - mlen - 1] = 0x00;
	memcpy(em.octets + (K_BYTES - mlen), M, sizeof(M[0]) * mlen);

	os2ip(&m, &em);
	rsaep(&c, n, e, &m);
	i2osp(C, &c);
	return 0;
}

int rsaes_pkcs1_v15_decrypt(uint8_t **M, size_t *mlen, const big_integer *n, const big_integer *d, const octet_string *C) {
	static octet_string em;
	static big_integer m, c;
	os2ip(&c, C);
	rsadp(&m, n, d, &c);
	i2osp(&em, &m);

	if (em.octets[0] != 0x00 || em.octets[1] != 0x02) {
		return 1;
	}
	int i;
	for (i = 2; i < K_BYTES; ++i) {
		if (em.octets[i] == 0x00) {
			break;
		}
	}
	if (i - 2 < 8 || i == K_BYTES) {
		return 1;
	}
	++i;

	*mlen = 128 - i;
	*M = (uint8_t*)malloc(sizeof(uint8_t) * (*mlen));
	memcpy(*M, em.octets + i, sizeof(uint8_t) * (*mlen));
	return 0;
}
